# -*- coding: utf-8 -*-
"""
-------------------------------------------------
   File Name：     common
   Description :
   Author :       Flyoung
   date：          2023/9/5
-------------------------------------------------
   Change Activity:
                   2023/9/5:
-------------------------------------------------
"""
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from functools import wraps
from db.config import db_url

engine = create_engine(db_url, echo=False, pool_size=2000)  # 建立引擎
Base = declarative_base()


def get_session() -> Session:
    """
    获取 session
    :return:
    """
    session = sessionmaker(bind=engine)
    return session()


def with_session(func):
    """
    查询时，自动获取 session
    :param func:
    :return:
    """

    @wraps(func)
    def wrapped(*args, **kwargs):
        session = get_session()
        try:
            result = func(session, *args, **kwargs)
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()
        return result

    return wrapped
